#!/usr/bin/python
# -*- coding:utf-8 -*-

"""
对外API 任务与任务结果管理
"""
import re
import arrow
import bson
from flask import g

from app.api.v3.task_util import CreateTask
from app.db import Pagination
from app.db.models.asset_tasks import AssetTasks
from app.libs.redprint import RedPrint
from app.libs.regexp import NAME_RE
from app.libs.utility import api_v3_response, get_request_args_v3
from app.db.models.tasks import Task
from app.db.models.balances import Balances
from app.db.models.jobs import Job
from app.db.models.wass import ScanEvents
from app.libs.utils import enable_task_group, balance_task_count_sync
from app.libs.enums import AVAILABLE_TASK_GROUP_TUPLE
from app.validators.asset import AssetStatus
from app.log.logs import CAssetTaskLog

api = RedPrint('task')


@api.route('/detail', methods=['GET'])
def detail():
    """某个任务的详情"""
    args = get_request_args_v3(["task_id", ])
    task_id = args.get("task_id")
    if not task_id:
        return api_v3_response(status="fail", errcode=2040001, data={})
    try:
        task_id = bson.ObjectId(task_id)
    except bson.errors.BSONError:
        return api_v3_response(status="fail", errcode=2040002, data={})
    allow_task_keys = ('id', 'taskSessionId', 'taskId', 'jobId', 'startTime', 'endTime', 'taskType',
                       'status', 'uid', 'progress', 'name', 'nextExecDate', 'triggerType')
    allow_job_keys = ('id', 'sourceIp', 'targetUrl')
    allow_asset_task_keys = ('id', 'taskGroupType', 'httpSettings', 'pingSettings')
    allow_event_keys = ('event_id', 'event_name', 'detail', 'severity', 'job_name')
    task = Task.objects.filter(uid=g.user.id, id=task_id, status="completed").only(*allow_task_keys).first()
    if not task:
        return api_v3_response(status="fail", errcode=2040003, data={})
    job = Job.objects.filter(id=task.jobId).only(*allow_job_keys).first()
    asset_task = AssetTasks.objects.find({f"{task.taskType}Settings.taskId": task.taskId}).only(*allow_asset_task_keys).first()
    if not all([job, asset_task]):
        return api_v3_response(status="fail", errcode=2040003, data={})
    data = {
        "task_id": str(task.id),
        "task_status": task.status,
        "progress": task.progress,
        "target": job.targetUrl,
        "target_ip": job.sourceIp or "",
        "task_type": task.taskType,
        "task_name": task.name or "",
        "target_id": str(job.id),
        "start_time": task.startTime,
        "task_session_id": task.taskSessionId,
        "end_time": task.endTime,
        "task_group_id": str(asset_task.id),
        "task_group_type": asset_task.taskGroupType,
        "next_exec_date": task.nextExecDate or "",
        "trigger_type": task.triggerType,
        "target_url": ""
    }
    if task.taskType in AVAILABLE_TASK_GROUP_TUPLE:
        data["target_url"] = getattr(asset_task, f"{task.taskType}Settings").target or ""

    event_list = []
    source_events = ScanEvents._get_collection().find({"task_session_id": task.taskSessionId}, {k: 1 for k in allow_event_keys})
    for event in source_events:
        if event["event_name"] == "statistics" and event["job_name"] == "appscan":
            continue
        if event["event_name"] == "site_info" and event["job_name"] == "kscan":
            continue
        event.pop("job_name", "")
        event.pop("_id", "")
        event_list.append(event)
    data["events"] = event_list
    return api_v3_response(data=data)


@api.route('/group/list', methods=['GET'])
def get_group_list():
    """任务组列表"""
    query_fields = ["limit", "page", "created_start_at", "created_end_at",
                    "target", "enable", "type", "name", "trigger_type"]
    request_args = get_request_args_v3(query_fields)
    typ = request_args.get("type")
    if (not typ) or typ not in ['host', 'web_security', 'web_asset', 'http', 'ping', 'ipv6']:
        return api_v3_response(status="fail", errcode=2040007, data={})
    query = {"uid": g.user.id, "taskGroupType": typ}
    enable = request_args.get("enable", None)
    true_list = ['true', True, "True"]
    false_list = [False, 'false', "False"]
    if (enable not in [None, ""]) and (enable not in true_list + false_list):
        return api_v3_response(status="fail", errcode=2040008, data={})
    elif (enable not in [None, ""]) and (enable in true_list):
        query["enableMonitor"] = True
    elif (enable not in [None, ""]) and (enable in false_list):
        query["enableMonitor"] = False
    if name := request_args.get("name"):
        query.update({"name": {"$regex": name}})
    if target := request_args.get("target"):
        query.update({"targetUrl": {"$regex": target}})
    trigger_type = request_args.get("trigger_type")
    if trigger_type and trigger_type not in ['schedule', 'manual']:
        return api_v3_response(status="fail", errcode=2040009, data={})
    elif trigger_type and trigger_type in ['schedule', 'manual']:
        query.update({"triggerType": trigger_type})
    start_at = request_args.get("created_start_at")
    end_at = request_args.get("created_end_at")
    if start_at:
        try:
            start_at = arrow.get(start_at).datetime
        except arrow.parser.ParserError:
            return api_v3_response(status="fail", errcode=2040010, data={})
    if end_at:
        try:
            end_at = arrow.get(end_at).datetime
        except arrow.parser.ParserError:
            return api_v3_response(status="fail", errcode=2040010, data={})
    if start_at:
        query.setdefault('createTime', {})['$gte'] = start_at
    if end_at:
        query.setdefault('createTime', {})['$lte'] = end_at
    all_ats = AssetTasks.objects.find(query).order_by("-id")
    page, limit = request_args.get("page", 1),  request_args.get("limit", 10)
    count = all_ats.count()
    data = []
    for item in Pagination(all_ats, page=page, per_page=limit, not_err=True).items:
        tmp = {
            "create_time": item.createTime,
            "target": item.targetUrl,
            "task_group_id": item.id,
            "enable": item.enableMonitor,
            "target_id": item.jobId,
            "task_group_type": item.taskGroupType,
            "trigger_type": item.triggerType,
            "target_url": "",
            "name": item.name
        }
        task_group_type = item.taskGroupType
        if task_group_type in AVAILABLE_TASK_GROUP_TUPLE:
            tmp["target_url"] = getattr(item, f"{task_group_type}Settings").target
        data.append(tmp)
    results = {"count": count, "page": page, "limit": limit, "results": data}
    return api_v3_response(data=results)


@api.route('/group/enable', methods=['POST'])
def update_task_group_enable():
    """开关周期任务"""
    request_args = get_request_args_v3(["task_group_id", "enable"])
    task_group_ids = request_args.get("task_group_id")
    enable = request_args.get("enable")
    if (not task_group_ids) or not isinstance(task_group_ids, list):
        return api_v3_response(status="fail", errcode=2040011, data={})
    id_count = len(task_group_ids)
    if id_count > 50:
        return api_v3_response(status="fail", errcode=2040012, data={})
    try:
        task_group_ids = [bson.ObjectId(t) for t in task_group_ids if t]
    except bson.errors.BSONError:
        return api_v3_response(status="fail", errcode=2040011, data={})
    if len(set(task_group_ids)) != id_count:
        return api_v3_response(status="fail", errcode=2040011, data={})
    if enable not in [True, False]:
        return api_v3_response(status="fail", errcode=2040008, data={})
    # 开启任务时， 判断一下资产是否已经认证
    if enable:
        job_ids = AssetTasks.objects.distinct(
            "jobId", {"uid": g.user.id, "_id": {"$in": task_group_ids}, "enableMonitor": not enable})
        affects = []
        for job in Job.objects.find({"_id": {"$in": job_ids}, "verification.isVerified": False}):
            affects.append(job.targetUrl)
        if affects:
            return api_v3_response(
                status="fail", errcode=2040020, data={}, message_kwargs={"affects": "、".join(affects)})
    data = []
    for at in AssetTasks.objects.find({"uid": g.user.id, "_id": {"$in": task_group_ids}, "enableMonitor": not enable}):
        at.update(enableMonitor=enable)
        enable_task_group(at)
        data.append(at.id)
    return api_v3_response(message=f'操作成功', data={"task_group_id": data})


@api.route('/monitor/status-update', methods=['POST'])
def monitor_status():
    """修改任务监测开关"""
    args = get_request_args_v3(["target_id", "enable_monitor", "type"])
    target_id = args.get("target_id")
    target_group_type = args.get("type")
    enable_monitor = args.get("enable_monitor")
    if enable_monitor is None or not isinstance(enable_monitor, bool):
        return api_v3_response(status="fail", errcode=2040004, data={})
    try:
        update_data = AssetStatus(**{"jobId": target_id}).dict()
    except Exception:
        return api_v3_response(status="fail", errcode=1040001, data={})

    if target_group_type and target_group_type not in ("available", "web", "host"):
        return api_v3_response(status="fail", errcode=2040005, data={})

    if target_group_type == "available":
        target_group_types = ["ping", "http"]
    elif target_group_type == "web":
        target_group_types = ["web_security"]
    elif target_group_type == "host":
        target_group_types = ["host"]
    else:
        target_group_types = ["ping", "http", "web_security", "host"]
    asset_tasks = AssetTasks.objects.filter(uid=g.audit_uid, jobId=update_data["jobId"], triggerType="schedule",
                                            taskGroupType__in=target_group_types)
    # 下发任务
    for asset_task in asset_tasks:
        if asset_task.enableMonitor == enable_monitor:
            continue
        asset_task.update(enableMonitor=enable_monitor)
        enable_task_group(asset_task=asset_task, insert=False, immediate_exec=False)
    balance = Balances.objects.filter(uid=g.user.id).first()
    if balance:
        balance_task_count_sync(balance.uid, balance)
    return api_v3_response(message=f'操作成功', data={})


@api.route('/create', methods=['POST'])
def create_task():
    """下发任务"""
    args = get_request_args_v3(["targets", "task_group_type", "name", "task_settings", "trigger_type"])
    targets = args.get("targets")
    task_group_type = args.get("task_group_type")
    task_settings = args.get("task_settings")
    name = str(args.get("name", ""))
    trigger_type = args.get("trigger_type")

    params_is_good = True
    error_code = 0
    balance = Balances.objects.filter(uid=g.user.id, planExpireTime__gte=arrow.utcnow().datetime).first()
    if task_group_type not in ["web_security", "http", "ping"]:
        error_code = 2040016
        params_is_good = False
    elif not CreateTask.check_balance_feature(balance, task_group_type):
        error_code = 2040019
        params_is_good = False
    elif not isinstance(targets, list):
        error_code = 2040013
        params_is_good = False
    elif len(targets) > 50 or len(targets) < 1:
        error_code = 2040015
        params_is_good = False
    elif trigger_type not in ["manual", "schedule", "scheduleIncrement"]:
        error_code = 2040017
        params_is_good = False
    elif task_group_type != "web_security" and trigger_type == "manual":
        error_code = 2040014
        params_is_good = False
    elif not NAME_RE.fullmatch(name):
        error_code = 1040019
        params_is_good = False
    elif task_group_type == "web_security" and not CreateTask.check_task_settings(task_settings):
        error_code = 2040018
        params_is_good = False
    if not params_is_good:
        return api_v3_response(status="fail", errcode=error_code, data={})

    fail_tasks = []
    success_tasks = []
    exist_tasks = []
    affects = []
    for target in targets:
        flag, msg = CreateTask(
            user=g.user,
            target_url=target,
            settings=task_settings,
            task_group_type=task_group_type,
            trigger_type=trigger_type
        ).create_task(name)
        if flag:
            success_tasks.append(target)
            affects.append(msg)
        elif (not flag) and re.findall(f"已经存在.*任务", str(msg)):
            exist_tasks.append(target)
        elif not flag:
            fail_tasks.append(target)

    if affects or (fail_tasks and "web_security" == task_group_type):
        log_str = "web安全监测添加单次任务" if trigger_type == "manual" else "web安全监测添加周期任务"
        log_info = {
            'affects': affects,
            'actionMessage': f'web安全监测',
            'detail': f'{log_str}: 任务名称：{name}, 添加任务成功资产地址: {",".join(affects)}, 添加任务失败资产地址: {",".join(fail_tasks)}'
        }
        CAssetTaskLog(**log_info).info(log_info['actionMessage'], log_info['detail'])
    if affects or (fail_tasks and task_group_type in ["http", "ping"]):
        log_info = {
            'affects': affects,
            'actionMessage': f'可用性监测',
            'detail': f'新建可用性监测任务: 添加任务成功资产地址: {",".join(affects)}, 添加任务失败资产地址: {",".join(fail_tasks)}'
        }
        CAssetTaskLog(**log_info).info(log_info['actionMessage'], log_info['detail'])
    return api_v3_response(message=f'下发任务, 成功{len(success_tasks)}个, 失败{len(fail_tasks)}个, 已存在{len(exist_tasks)}个',
                           data={"fail": fail_tasks, "success": success_tasks, "exist": exist_tasks})
